import cv2
import numpy as np

from utils import export_results


def is_black_frame(frame, threshold=30):
    '''
    判断是否为黑屏
    黑屏亮度阈值（0-255）
    :param frame:
    :param threshold:
    :return:
    '''
    # 计算帧的平均值
    avg_color = np.mean(frame)
    # 如果平均值小于阈值，则认为是黑屏
    return avg_color < threshold


def check_black_frames(video_path, threshold=30):
    cap = cv2.VideoCapture(video_path)  # 打开视频文件
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 获取视频总帧数
    fps = cap.get(cv2.CAP_PROP_FPS)  # 获取视频的帧率
    sample_interval = total_frames // 20  # 计算每20等分的间隔

    black_frames = []  # 用于存储黑屏帧的信息

    for i in range(0, total_frames, sample_interval):  # 遍历每个采样帧
        cap.set(cv2.CAP_PROP_POS_FRAMES, i)  # 设置当前帧的位置
        ret, frame = cap.read()  # 读取当前帧
        if not ret:  # 如果帧读取失败，跳出循环
            break
        if is_black_frame(frame, threshold):  # 检查当前帧是否为黑屏
            timestamp = i / fps  # 计算时间戳（秒）
            print(f"发现黑屏帧在第 {i} 帧，时间戳：{timestamp:.2f} 秒")  # 输出黑屏帧的帧数和时间戳
            black_frames.append({'frame': i, 'timestamp': timestamp})  # 将黑屏帧信息添加到列表中

    cap.release()  # 释放视频捕获对象
    return black_frames  # 返回黑屏帧的信息列表


def process_videos_black(video_paths, threshold=30, output_format='csv', output_file='black_frames_results'):
    '''
    批量处理视频文件
    :param video_paths:
    :param threshold:
    :param output_format:
    :param output_file:
    :return:
    '''
    all_black_frames = {}  # 用于存储所有视频的黑屏帧信息
    for video_path in video_paths:  # 遍历每个视频文件
        print(f"处理视频: {video_path}")
        black_frames = check_black_frames(video_path, threshold)  # 检查黑屏帧
        all_black_frames[video_path] = black_frames  # 存储结果

    export_results(all_black_frames, output_format, output_file)  # 导出汇总结果


if __name__ == '__main__':
    path1 = '/Users/maweidong/source/video_audio_inspector/7M文件大小.mp4'
    path2 = '/Users/maweidong/source/video_audio_inspector/18M大小文件.mp4'
    video_file_paths = [path1, path2]  # 替换为你的视频文件路径列表
    user_threshold = 30  # 用户自定义的黑屏亮度阈值，可以根据需要修改
    output_format = 'csv'  # 用户选择的输出格式，可以是 'csv' 或 'json'

    # 调用函数处理视频
    process_videos_black(video_file_paths, user_threshold, output_format)
